# !/usr/bin/env python
# encoding: utf-8


"""
  @author: gaogao
  @file: db_helper.py
  @time: 2020/12/10 15:59
  @desc:
"""

from sqlalchemy.orm import sessionmaker, scoped_session
from database.base import Base
from contextlib import contextmanager
from sqlalchemy import create_engine
import functools

# # 数据库连接配置
MYSQL = {
    'db_macrochip': {
        'user': 'root',
        'passwd': 'admin',
        'host': '127.0.0.1',
        'port': 3306,
        'db': 'macrochip',
    }
}

engines = {}


def init_engines():
    """初始化数据库连接"""
    for k, v in MYSQL.iteritems():
        mysql_url = ("mysql+pymysql://{user}:{passwd}@{host}:{port}/{db}"
                     "?charset=utf8".format(**v))
        engines[k] = create_engine(mysql_url,
                                   pool_size=10,
                                   max_overflow=-1,
                                   pool_recycle=1000,
                                   echo=False)


# 初始化所有数据库的连接，后续如果新增数据库访问，在MYSQL里面直接加入数据库配置即可
init_engines()


class DBHelper:

    # def __init__(self, model):
    #     self.model = model
    #     self.model.metadata.create_all(Base.engine)

    def create_table(self, model):
        self.model = model
        self.model.metadata.create_all(Base.engine)

    def get_session(self):
        """获取session"""
        return scoped_session(sessionmaker(bind=Base.engine, expire_on_commit=False))()

    # 定义上下文函数，便能够自动进行事务管理
    @contextmanager
    def db_session(self):
        session = self.get_session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            if session:
                session.close()

    # 添加一条数据
    def insert_data(self, entity):
        with self.db_session() as session:
            session.add(entity)

    # 添加多条数据
    def insert_batch_data(self, entity_list):
        with self.db_session() as session:
            session.add_all(entity_list)

    # 删除数据
    def delete_data(self):
        with self.db_session() as session:
            session.query(self.model).filter().delete(synchronize_session=False)

    # 修改数据
    def update_data(self):
        with self.db_session() as session:
            session.query(self.model).filter().update()

    # 查询一条结果
    def query_one(self):
        with self.db_session() as session:
            return session.query(self.model).filter().one()

    # 查询多条
    def query_all(self):
        with self.db_session() as session:
            return session.query(self.model).filter().all()


def class_dbsession(commit=True):
    """用于BaseModel中进行数据库操作前获取dbsession操作.
    :param commit:进行数据库操作后是否进行commit操作的标志，True：commit, False:不commit
    """
    def wrapper(func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            # cls为对象或类
            cls = args[0]
            # 实际传入的参数
            new_args = args[1:]
            with db_session() as session:
                return func(cls, session, *new_args, **kwargs)
        return inner
    return wrapper


# 定义上下文函数，便能够自动进行事务管理
@contextmanager
def db_session(db='macrochip', commit=True):
    session = get_session(db)
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        raise e
    finally:
        if session:
            session.close()


def get_session(db):
    """获取session"""
    return scoped_session(sessionmaker(bind=engines[db], expire_on_commit=False))()
